Decompose Wannier-function densities onto a Gaussian × spherical-harmonic basis by elinscott · Pull Request #648 · wannier-developers/wannier90

elinscott · 2026-07-06T16:35:00Z

This PR is a WIP. I'm very happy to discuss the big picture but there's no need to start reviewing the code details just yet

Summary

This PR adds an optional post-processing step to wannier90.x that decomposes each Wannier function's density $|w_n(r)|^2$ onto an orthonormalised basis of Gaussian radial functions × real spherical harmonics centred on the WF's own centre, writing per-WF coefficient and power-spectrum files. These SOAP-style descriptors follow Himanen et al., Adv. Sci. 7, 1902333 (2020).

The motivating use case is machine-learning workflows that need compact, rotation-covariant descriptors of individual Wannier functions. The koopmans package uses these power spectra to predict screening parameters (see Schubert et al., npj Comput. Mater. 10 299 (2024)).

Because the decomposition is very general and can be performed purely on wannier90 content, I thought it would naturally slot into wannier90. I realise it is a substantial addition and am very happy to discuss further or publish it is a separate code if you think it is genuinely beyond the scope of what you want Wannier90 to do.

Theory

Each Wannier density $\rho_n(\mathbf{r}) = |w_n(\mathbf{r})|^{2}$ is expanded about its centre $\mathbf{c}_n$ as

$$c^{(n)}_{n'lm} = \int_{|\mathbf{r}-\mathbf{c}_n| \le r_{\mathrm{cut}}} \rho_n(\mathbf{r}) g_{n'l}(|\mathbf{r}-\mathbf{c}_n|) Y_{lm}(\theta,\varphi) \mathrm{d}V,$$

evaluated in terms of radial $g_{nl}$ and angular $Y_{lm}$ basis functions over the Born–von-Kármán grid. The rotationally invariant power spectrum is then given by

$$p^{(n)}_{n_1 n_2 l} = \sum_{m=-l}^{l} c^{(n)}_{n_1 l m} c^{(n)}_{n_2 l m}, \qquad n_1 \le n_2,$$

Radial basis

The primitive radial functions are Gaussians,

$$\phi_{nl}(r) = r^{l} e^{-\alpha_{nl} r^{2}}, \qquad \alpha_{nl} = -\frac{\ln(\mathrm{thr}/r_{\mathrm{thr},n}^{l})}{r_{\mathrm{thr},n}^{2}},$$

with $\mathrm{thr} = 10^{-3}$ and the threshold radii $r_{\mathrm{thr},n}$ equally spaced on $[r_{\mathrm{min}}, r_{\mathrm{max}}]$ (decompose_r_min, decompose_r_max), i.e. each Gaussian has decayed to the threshold at its $r_{\mathrm{thr},n}$. Their overlaps are evaluated analytically,

$$S_{nn'l} = \int_0^{\infty} \phi_{nl}(r) \phi_{n'l}(r) r^{2} \mathrm{d}r = \frac{1}{2} \Gamma(l + 3/2) (\alpha_{nl}+\alpha_{n'l})^{-(l+3/2)},$$

and the orthonormal radial basis is obtained by Löwdin symmetric orthogonalisation,

$$g_{nl}(r) = \sum_{n'} \beta_{n'nl} \phi_{n'l}(r), \qquad \beta^{(l)} = V \mathrm{diag}(e^{-1/2}) V^{T},$$

where $V$, $e$ are the eigenvectors/eigenvalues of $S^{(l)}$.

Angular basis

The angular part uses real spherical harmonics $Y_{lm}(\theta,\varphi)$ in the scipy real convention (Condon–Shortley phase, $\sqrt{2}$-normalised $m \neq 0$ components).

What the code does

With wannier_decompose = .true. (runs during the plotting phase, also under restart = plot):

Builds each WF on the real-space grid of the Born–von-Kármán supercell (the mp_grid multiple of the unit cell).
Normalises each |w_n|² to unit integral.
Projects each selected WF density onto the basis functions within a sphere of radius decompose_r_cut about the WF centre, using the minimum-image convention
Writes, per WF n:
- <seed>_NNNNN.coeff — expansion coefficients c_nlm
- <seed>_NNNNN.power — the invariant power spectrum p(n₁,n₂,l) built from those coefficients.
Optionally, if decompose_centres_file is given: also decomposes the group density (sum of all this run's normalised WF densities) about every centre listed in that file, writing <seed>_gc_NNNNN.coeff.

New keywords

Keyword	Type	Default	Meaning
`wannier_decompose`	logical	`.false.`	Enable the decomposition
`decompose_r_cut`	real (Å)	— (required)	Cutoff radius of the projection sphere
`decompose_n_max`	integer	6	Number of radial basis functions per l
`decompose_l_max`	integer	6	Maximum angular momentum
`decompose_r_min`	real (Å)	0.5	Smallest Gaussian decay radius
`decompose_r_max`	real (Å)	4.0	Largest Gaussian decay radius
`decompose_list`	list	all WFs	Which WFs get own-centre output (group density always sums all WFs)
`decompose_centres_file`	string	—	File of extra centres (Å, Cartesian) for group-density decomposition

Implementation

src/decompose.F90 (new) — module w90_decompose: basis math (decompose_radial_params, decompose_real_ylm), projection (decompose_project), power spectra (decompose_power_orb), and the root-only driver + I/O (decompose_main)
src/plot.F90 — the WF real-space grid construction inside plot_wannier is extracted into a shared helper plot_build_wannier_grid
src/wannier90_readwrite.F90 / src/readwrite.F90 / src/wannier90_types.F90 — support for new keywords.

Validation and tests

Basis functions, projection coefficients and power spectra were validated against the legacy Python implementation in koopmans (coefficients agree to ~2·10⁻⁷ relative — limited by the reference's numerical-quadrature overlaps vs. our analytic ones; power spectra to ~7·10⁻¹⁵ on identical inputs).
New test-suite cases (serial, group plot):
- testw90_decompose, testw90_decompose_power, testw90_decompose_gc — GaAs (inputs shared with testw90_example01), covering own-centre coefficients, power spectra and the group-density/centres-file channel, with new parsers in test-suite/tools/parsers/;
- testw90_decompose_sphere — a synthetic spherical density in a triclinic cell)
- testw90_decompose_l2 — a synthetic pure-l=2 density in a cubic cell

Documentation

User guide (docs/docs/user_guide/wannier90/parameters.md), parameters.xml, plot-parameters CSV, and a CHANGELOG entry.

Limitations / notes

Root-only implementation: the decomposition itself runs on the root rank after the (already parallel) grid accumulation. Should be fine for the intended grid sizes; parallelisation is possible later if needed.
Spinor WFs are not supported for the moment.
The total-density-across-runs story assumes integer occupations (each run's group density is the sum of its normalised WF densities)
The UNK-file error-handling fix from Fix MPI hang/abort on UNK read failure in plot_wannier #647 is merged into this branch, applied inside the extracted plot_build_wannier_grid helper (where the k-point read loop now lives)

AI Declaration

I have used Claude extensively to port the legacy code (originally written in python by a student of mine) over to fortran, document the code, etc. I am in the process of reviewing all the code. It has been validated against an independent reference implementation and the new test-suite cases described above.

New module implementing the basis mathematics for decomposing Wannier function densities onto Lowdin-orthonormalised Gaussian radial functions times real spherical harmonics (Himanen et al., Adv. Sci. 7, 1902333 (2020)): - decompose_radial_params: analytic Gaussian decay coefficients and Lowdin orthonormalisation (S^-1/2 via LAPACK dsyev) of the analytic radial overlap matrix. - decompose_real_ylm: real spherical harmonics for all l <= l_max via the standard associated-Legendre upward recurrence, scipy real-Ylm convention (Condon-Shortley phase, sqrt(2)*(-1)^m real combinations). Density projection, power spectra, input keywords and file I/O follow in subsequent commits. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Pure refactor of plot_wannier: the real-space Wannier-function grid build (per-k UNK read, disentanglement rotation, band contraction with the U matrices, Bloch-phase accumulation, comms_reduce) moves into a new public subroutine plot_build_wannier_grid, parameterized over the WF index list and the supercell grid bounds so that other consumers can build WF grids on different supercells without going through wannier_plot settings. plot_wannier now calls the helper and keeps the spinor combination, global-phase fixing, reality check and xsf/cube output unchanged. No behavior change: testw90_cube_format passes and its cube output is byte-identical to a pre-refactor build (modulo the embedded run timestamp); numerics, loop order and MPI structure are untouched. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Extend the w90_decompose module with the decomposition core: - decompose_project: project a real-space density onto the orthonormalised Gaussian-radial x real-Ylm basis about a given centre, using minimum-image Cartesian displacements (general, non-orthorhombic cells) and a spherical mask r <= r_cut, integrated as a plain sum with the exact cell volume element. - decompose_power_orb: rotationally invariant power spectrum p_{n1 n2 l} = sum_m c_{n1 l m} c_{n2 l m} (n2 >= n1). - decompose_main: root-only driver taking the WF grid from plot_build_wannier_grid on the Born-von-Karman supercell; normalises each |w_n|^2 to unit integral, decomposes each selected WF about its own centre, and optionally decomposes the group density (sum of all WF densities) about a list of externally supplied centres. Writes self-describing ASCII <seed>_NNNNN.coeff/.power and <seed>_gc_NNNNN.coeff files. - Validates r_cut against the inscribed-sphere radius of the BvK supercell, and propagates radial-basis failures (dsyev, non-positive overlap eigenvalues) through w90_error. Validated against the reference python implementation (koopmans ml module): projection agrees to ~4e-13 on identical domains, coefficient fixtures reproduced to ~2e-7 (limited by quadrature-vs-analytic radial overlaps), power spectra to ~7e-15. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

New top-level keyword wannier_decompose with settings decompose_n_max (default 6), decompose_l_max (6), decompose_r_min (0.5 A), decompose_r_max (4.0 A), decompose_r_cut (required), decompose_list (default: all WFs) and decompose_centres_file (optional group-density channel): - decompose_type added to wannier90_types.F90 and lib_common_type; parsed by a new reader in wannier90_readwrite.F90 following the wannier_plot patterns; keywords registered in the master keyword list in readwrite.F90; settings echoed in the PLOTTING summary block. - plot_main gains a decompose branch: plot_decompose builds the WF grid on the Born-von-Karman supercell (mp_grid) via the shared plot_build_wannier_grid and hands it to decompose_main. Runs under restart = plot. Spinor wavefunctions are rejected with a clear error. - decompose_main gains an explicit output-subset argument so that decompose_list only selects which WFs are decomposed about their own centres, while the group density always sums over all of the run's Wannier functions (its coefficients are combined across runs, so it must not depend on the output selection). - The external centres file is ASCII: one Cartesian centre (Angstrom) per line, '#' comments allowed. Smoke-tested on the GaAs cube-format test inputs: coefficient/power files produced for own-centre and group-density channels; r_cut validation reports the computed Born-von-Karman bound; restricting decompose_list leaves the group-density output byte-identical. testw90_cube_format passes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

User-guide prose for all eight new keywords (basis definition, output file formats and orderings, the Born-von-Karman r_cut bound, the centres-file format, and the cross-run group-density summation use case), parameters.xml entries and the corresponding rows in the generated plot-parameters table, and a CHANGELOG entry. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Three GaAs tests (inputs symlinked from testw90_example01, following the example03 sibling-directory pattern since testcode checks a single output file per test): testw90_decompose checks the own-centre coefficients of WF 1, testw90_decompose_power the orbital-orbital power spectrum, and testw90_decompose_gc the group-density coefficients about an external centre supplied via decompose_centres_file. New parsers for the .coeff/.power formats; coefficient values compared at abs 1e-8 / rel 1e-6. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

decompose_project placed grid array index i at fractional coordinate i/ng (a zero-based convention, index 0 at the cell origin), but the Wannier-function grid built by plot_build_wannier_grid puts index i at (i-1)/ng (index 1 at the origin, matching its Bloch-phase factors). On the production path every density sample was therefore displaced by one grid cell along each lattice vector relative to the WF centre -- about 0.5 Angstrom for the GaAs test case -- manufacturing spurious l >= 1 content in the coefficients and biasing the power spectra. The original unit validation did not catch this because its driver passed zero-based array bounds, for which the old formula was self-consistent; the mismatch only appears against the one-based grid of plot_build_wannier_grid. Verified both ways: a synthetic spherically-symmetric density about a non-grid-aligned centre on a one-based grid drops from an l>=1/l=0 coefficient ratio of 3.7e-1 to 3.9e-6 (pure grid discretisation residue) with this fix, and the fixed routine fed via one-based bounds still reproduces the reference python implementation's coefficients to 2.1e-7 (unchanged). The three decompose test benchmarks are regenerated accordingly. Found by independent review. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Rename locals that shadowed an intrinsic (exponent) or misdescribed their contents (norm holding the Condon-Shortley sign, fname holding an error message); derive the supercell volume from the face cross-products instead of a hand-expanded triple product; hoist the inverse-lattice and volume-element computation out of decompose_project into its caller; document that the shifted theta/phi convention is inherited from the reference implementation deliberately. No behavior change: all decompose outputs verified byte-identical against the previous binary on the GaAs smoke inputs; decompose and cube-format tests pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Two synthetic gamma-only single-WF tests whose coefficients are known analytically, so the benchmarks pin correctness rather than just reproducibility: - testw90_decompose_sphere: a spherical Gaussian orbital about an off-grid centre; all l >= 1 coefficients must vanish (measured l>=1/l=0 ratio 9e-12; the grid-origin bug fixed in the previous commits gives 0.4, so this test fails hard against it), and the l = 0 coefficients match analytic radial quadrature to 8e-11. - testw90_decompose_l2: a d_z2-like orbital; the density contains only (l,m) = (0,0), (2,0), (4,0), with the six allowed coefficients matching the analytic Legendre decomposition of (3t^2-1)^2 to 5e-7 and odd-l / m /= 0 content below 1e-7. Inputs (UNK, .amn, .mmn) are generated deterministically by the committed generator script; the .mmn overlaps are computed from the grid so the Wannier centres W90 reports equal the true orbital centres to all printed digits. A checker script documents the l-purity and analytic comparisons. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The spherical physics-pinning test now runs in a deliberately irregular triclinic cell (7 x 8 x 9 Angstrom, angles 70/80/75 degrees) instead of a cubic one, so the analytically-known benchmark also exercises the non-orthorhombic geometry path: general fractional/Cartesian transforms, minimum-image over a skewed Born-von-Karman cell, and the general volume element. A cubic cell cannot expose lattice-convention errors -- its lattice matrix equals its transpose -- whereas here re-projecting the same density with a transposed lattice degrades the l-purity from 1.6e-8 to 5.0e-2, so the test discriminates. Measured with the correct conventions: l>=1 to l=0 coefficient ratio 1.6e-8, l=0 coefficients matching the analytic radial quadrature to ~1e-8, and the W90-reported Wannier centre matching the true orbital centre to 2.9e-7 Angstrom up to a lattice translation (the b.c > pi phase wrap for the skewed cell; the checker compares minimum-image). The generator now accepts a general 3x3 cell and takes its b-vectors from wannier90.x -pp (12 vectors in 6 shells here, not the cubic 6); the cubic d_z2 test regenerates bit-identically. r_cut drops from 3.8 to 3.0 Angstrom to respect the triclinic inscribed-sphere bound (3.366 Angstrom). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Symptom: with MPI and wannier_plot enabled, a missing, unreadable, or header-mismatched UNK file for one k-point could hang or abort the run with no clean error message, instead of reporting the offending file. Cause: k-points are distributed over ranks (dist_k), so each rank reads only its own UNK files. Two rank-local failure sites inside the read loop were mishandled: 1. A header mismatch called set_error_file directly. That routine syncs the error across the communicator (comms_sync_error -> mpi_allreduce). When only the owning rank failed, it entered the allreduce while the other ranks proceeded to the comms_reduce of wann_func: mismatched collectives, i.e. a deadlock (observed as an abnormal termination on OpenMPI, an indefinite hang on stricter MPI stacks). 2. A missing file was not detected: open without status='old' recreated it empty and the subsequent unguarded read aborted with a raw EOF runtime error and no W90 message. Fix: detect missing files (inquire + status='old'), add iostat to the open and all UNK reads, and record any rank-local failure in a local status flag instead of erroring mid-loop. After the loop every rank synchronises together: failing ranks call set_error_file (naming the file), succeeding ranks call comms_sync_error(comm, error, 0), so the collective is matched and all ranks return before comms_reduce. This is the same reduce-then-set idiom already used elsewhere (e.g. the unlucky_rank path in wannier_prog.F90 and write_kmesh in library_extra.F90). Repro (test-suite testw90_cube_format, 8 k-points, formatted UNK): mpirun -np 2 wannier90.x gaas # rank 1 owns k=5-8 deleting or corrupting UNK00005.1 previously hung/aborted; it now exits nonzero within seconds with "plot_wannier: file UNK00005.1 not found (rank: 1)". Control runs (all files present) are unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The UNK-reading k-point loop the fix patches inside plot_wannier was moved to plot_build_wannier_grid on this branch, so the conflict is resolved by applying the fix (status='old' + iostat on all UNK opens and reads; rank-local failure flag, then collective error raise / comms_sync_error after the loop, before the comms_reduce) inside the helper, which now serves both plot_wannier and plot_decompose. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The make-based build (src/obj/Makefile) needs the new module listed in OBJS, its dependency rule, and a plot.o prerequisite, otherwise plot.F90 fails with 'Cannot open module file w90_decompose.mod'. The cmake build already had it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

jaemolihm · 2026-07-06T17:05:09Z

Hello @elinscott , this looks very interesting! I have some questions.

In the decomposition, could you also decompose w_i(r) directly instead of decomposing rho_i(r) = |w_i(r)|^2? I think this would directly connect the atomic-orbital-like (s,p,d,...) nature of WFs.
How do you deal with the real-space grid? In plane-wave code like QE, wavefunctions are represented in plane waves, then taking overlap with spherical harmonics is much easier. In your Wannier implementation, do you also FFT to plane waves and compute there? Or do you take overlaps in real space? In the latter case, I guess there could be some aliasing error when the Wannier centers are not exactly on the FFT grid.
Practically, I often find storing the full UNK in a real-space grid is quite demanding for disk usage. So I prefer to postprocess using the plane-wave representation of the DFT code directly rather than converting to real-space UNK files, although this means the postprocessing can no longer be DFT-code-agnostic. I wonder what your experiences are.

elinscott · 2026-07-06T19:34:03Z

Hi @jaemolihm thanks for the nice questions! Let me answer them below

In the decomposition, could you also decompose w_i(r) directly instead of decomposing rho_i(r) = |w_i(r)|^2? I think this would directly connect the atomic-orbital-like (s,p,d,...) nature of WFs.

Sure, the machinery is agnostic. We went with density for a few reasons:

the motivating application comes from the world of orbital-density-dependent functionals, where the orbital density is our fundamental object
dealing with $|w_n|^2$ removes the arbitrary global sign/gauge of WFs (which would cause uniqueness issues when used as a descriptor for ML, though I'm sure there would be some way of constructing a clever descriptor that would digest it...)

How do you deal with the real-space grid? In plane-wave code like QE, wavefunctions are represented in plane waves, then taking overlap with spherical harmonics is much easier. In your Wannier implementation, do you also FFT to plane waves and compute there? Or do you take overlaps in real space? In the latter case, I guess there could be some aliasing error when the Wannier centers are not exactly on the FFT grid.

I do it naively: overlaps are taken in real space. WF centres don't lie on the grid, so yes, there is a discretisation error for off-grid centres. In practice it's small: the test suite includes a synthetic spherical density about a deliberately off-grid centre in a triclinic cell, and the spurious $l \ge 1$ content is $\sim 10^{-8}$ of the $l=0$ signal. (That said, that benchmark is deliberately smooth on the grid scale; for realistic densities with sharper radial basis functions the discretisation error would be larger... I will investigate more!)

I would absolutely consider a reciprocal-space grid implementation if it turned out there were situations where the aliasing causes problems.

Practically, I often find storing the full UNK in a real-space grid is quite demanding for disk usage. So I prefer to postprocess using the plane-wave representation of the DFT code directly rather than converting to real-space UNK files, although this means the postprocessing can no longer be DFT-code-agnostic. I wonder what your experiences are.

You're absolutely right, full real-space UNKs are heavy. This was a deliberate tradeoff to get the code sitting entirely within wannier90 rather than writing something DFT-code-specific. I'm testing it at the moment and thus far it hasn't been prohibitively cumbersome... but that experience may well change!

jaemolihm · 2026-07-06T21:43:25Z

Thanks a lot for the detailed reply! Looking forward to seeing how Gaussian the WFs are 😃

…se_radial_params For decompose_l_max >= 10 combined with decompose_r_min around 0.5 Angstrom the analytic decay condition yields a negative alpha (e.g. alpha = -0.0949 at l=10, r_thr=0.5), which makes the Gaussian overlap elements cc**(-l-3/2) NaN. The existing eigenvalue guard any(eval <= 0) cannot catch this because every comparison with NaN is false, so the failure mode depended on LAPACK dsyev happening to return info /= 0. Check the alphas explicitly (alpha <= 0 or NaN via the x /= x idiom) right after they are computed and raise a fatal error naming decompose_r_min/decompose_l_max as the knobs to adjust; also widen the local message buffer so the internal write cannot overflow. Defaults (l_max=6) are unaffected: all alphas stay positive and the synthetic decompose tests still pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

elinscott and others added 14 commits July 3, 2026 14:18

Appease pre-commit: fprettify alignment + markdownlint emphasis style

9ea376c

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

JeromeCCP9 added this to the v4.1 milestone Jul 17, 2026

Navigation Menu

Uh oh!

Decompose Wannier-function densities onto a Gaussian × spherical-harmonic basis - #648

Decompose Wannier-function densities onto a Gaussian × spherical-harmonic basis#648
elinscott wants to merge 15 commits into
wannier-developers:developfrom
elinscott:feature/decompose

elinscott commented Jul 6, 2026

Uh oh!

jaemolihm commented Jul 6, 2026 •

edited

Loading

Uh oh!

elinscott commented Jul 6, 2026 •

edited

Loading

Uh oh!

jaemolihm commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

elinscott commented Jul 6, 2026

Summary

Theory

Radial basis

Angular basis

What the code does

New keywords

Implementation

Validation and tests

Documentation

Limitations / notes

AI Declaration

Uh oh!

jaemolihm commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

elinscott commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jaemolihm commented Jul 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jaemolihm commented Jul 6, 2026 •

edited

Loading

elinscott commented Jul 6, 2026 •

edited

Loading